"""
卡尔曼滤波+布林带策略(zipline实现)
策略逻辑：
1. 使用卡尔曼滤波估计价格均衡水平
2. 计算布林带作为波动率通道
3. 当价格突破下轨且偏离均衡时买入
4. 当价格突破上轨且偏离均衡时卖出
5. 使用凯利公式进行头寸管理
"""

from zipline.api import order, record, symbol, set_commission
from zipline.finance.commission import PerShare
import numpy as np
from pykalman import KalmanFilter

def initialize(context):
    context.asset = symbol('AAPL')
    context.period = 20     # 布林带周期
    context.devfactor = 2.0  # 标准差倍数
    
    # 卡尔曼滤波参数
    context.transition_covariance = 0.01
    context.observation_covariance = 0.1
    
    # 设置佣金
    set_commission(PerShare(cost=0.001, min_trade_cost=1))

def handle_data(context, data):
    # 获取历史价格
    prices = data.history(context.asset, 'price', 100, '1d')
    if len(prices) < 100:
        return
    
    # 运行卡尔曼滤波
    kf = KalmanFilter(
        transition_matrices=[1],
        observation_matrices=[1],
        initial_state_mean=prices[0],
        initial_state_covariance=1,
        transition_covariance=context.transition_covariance,
        observation_covariance=context.observation_covariance
    )
    state_means, state_covariances = kf.filter(prices.values)
    predicted = state_means[-1][0]
    
    # 计算布林带
    recent_prices = prices[-context.period:]
    mean = np.mean(recent_prices)
    std = np.std(recent_prices)
    upper_band = mean + context.devfactor * std
    lower_band = mean - context.devfactor * std
    
    # 当前价格
    current_price = prices[-1]
    
    # 获取当前持仓
    current_position = context.portfolio.positions[context.asset].amount
    
    # 交易逻辑
    if not current_position:  # 没有持仓
        if current_price <= lower_band and current_price < predicted:
            order(context.asset, 100)  # 买入100股
    elif current_position > 0:  # 当前持有多头
        if current_price >= upper_band and current_price > predicted:
            order(context.asset, -current_position)  # 平仓
    
    # 记录指标
    record(price=current_price, predicted=predicted,
           upper_band=upper_band, lower_band=lower_band,
           position=current_position)

if __name__ == '__main__':
    import pandas as pd
    from zipline.utils.run_algo import run_algorithm
    # 回测配置
    start = pd.Timestamp('2020-01-01', tz='utc')
    end = pd.Timestamp('2023-01-01', tz='utc')
    result = run_algorithm(
        start=start,
        end=end,
        initialize=initialize,
        handle_data=handle_data,
        capital_base=100000,
        data_frequency='daily',
        bundle='quandl'
    )